import fitz  # PyMuPDF
import re


class PdfFitzAnalyse:
    def __init__(self, pdf_path, chapter_title_regex_list=None, page_num_prefix=None):
        self.pdf_path = pdf_path
        # 定义章节中的目录
        if chapter_title_regex_list is None:
            self.chapter_title_regex_list = [r"第 (\d+) 章 .+", r"第(\d+) 节 .+", r"(\d+\.\d+\.?\d*)\s.*"]
        if page_num_prefix is None:
            self.page_num_prefix = "_page_num_"

    def get_chapters_and_content(self):
        """
        以键值方式返回标题，和内容
        """
        chapters = {}
        current_chapter = None
        with fitz.open(self.pdf_path) as doc:
            # 遍历每一页
            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                text = page.get_text()
                for chapter_title_regex in self.chapter_title_regex_list:
                    # 查找章节标题
                    for match in re.finditer(chapter_title_regex, text):
                        current_title = match.group(0)
                        # 如果之前有章节，则记录之前章节的内容
                        if current_chapter:
                            chapter_text = self.replace_newline_after_subsection("".join(chapter_text))
                            chapters[current_chapter] = chapter_text
                        # 更新当前章节
                        current_chapter = current_title.replace("\n", " ") + f"{self.page_num_prefix}{page_num}"
                        chapter_text = [self.replace_text_line_change(text[match.end():])]  # 从匹配结束位置开始记录
                        break  # 跳出内层循环，避免重复匹配
                    # 继续记录当前章节的内容
                    if current_chapter:
                        chapter_text.append(self.replace_text_line_change(text))

                # 记录最后一个章节的内容
                if current_chapter:
                    text_content = "".join(line for line in chapter_text
                                           if not (self.is_table(line) or self.is_math_formula(line)))
                    text_content = self.replace_newline_after_subsection(text_content)
                    chapters[current_chapter] = text_content
        return chapters

    def get_page_num_prefix(self):
        return self.page_num_prefix

    @staticmethod
    def is_table(text):
        """
        简单判断表格的规则，例如包含多个行列的结构
        todo 待完善
        """
        if re.search(r'\|\s*\d+', text):
            return True
        return False

    @staticmethod
    def is_math_formula(text):
        """
        简单判断数学公式的规则，例如包含特殊字符或数学符号
        todo 待完善
        """
        math_symbols = r'[\+\-\*/=<>^√πΣ∫στ]'
        if re.search(math_symbols, text):
            return True
        return False

    @staticmethod
    def replace_text_line_change(text):
        """
        替换换行符
        """
        max_length = 65535
        lines = text.split("\n")
        text_result = ""
        for i, line in enumerate(lines):
            if line and line.endswith("。"):
                if len(text_result) > max_length:
                    text_sub = text_result[max_length:]
                    text_result = text_result[:max_length] + "\n" + text_sub
                if len(text_result) + len(line) > max_length:
                    text_result += "\n" + line
                    continue
                else:
                    text_result += line + "\n"
                    continue
            if i >= len(lines) - 1:
                continue
            if len(text_result) + len(line) > max_length:
                text_result += "\n" + line
            else:
                text_result += line
        return text_result

    @staticmethod
    def replace_newline_after_subsection(text):
        """
        替换小节标题后面的换行符
        """
        return re.sub(r'(\d+\.\d+(\.\d+)*)(\n)', r'\1 ', text)

